import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pydotplus
import seaborn as sb
from IPython.display import Image
from sklearn import tree
from sklearn.metrics import confusion_matrix
from data import DataUtils
import os

class VisualUtils(object):
	plt = plt

	def __init__(self):
		VisualUtils.plt.figure(figsize=(8, 5), dpi=100)
		VisualUtils.plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置显示中文字体
		VisualUtils.plt.rcParams['axes.unicode_minus'] = False  # 设置正常显示符号

	@staticmethod
	def plt_linear(test_y, pred_y, r2):
		plt.scatter(test_y, pred_y)
		plt.title('线性回归')
		plt.xlabel('Actual values')
		plt.ylabel('Predicted values')
		plt.plot(np.unique(test_y), np.poly1d(np.polyfit(test_y, pred_y, 1))(np.unique(test_y)))
		plt.text(3, 3, 'R-squared = %0.2f' % r2)
		plt.show()

	@staticmethod
	def plt_tree(model):
		# 要判断的文件路径
		file_path = "tree.pdf"
		if os.path.exists(file_path):
			return
		dot_data = tree.export_graphviz(model, out_file=None, feature_names=DataUtils.names[:8],
			filled=True, rounded=True, special_characters=True)
		graph = pydotplus.graph_from_dot_data(dot_data)
		Image(graph.write_pdf('tree.pdf'))

	@staticmethod
	def plt_pie(x_axis, y_axis, title):
		VisualUtils.plt.pie(x_axis, labels=y_axis, autopct="%1.2f%%", startangle=180)
		VisualUtils.plt.title(title)
		VisualUtils.plt.axis("equal")
		VisualUtils.plt.show()

	@staticmethod
	def plot_confusion_matrix(y_true, pred_y, title):
		# 压缩数据
		y_true = DataUtils.compress_data(pd.Series(y_true))
		pred_y = DataUtils.compress_data(pd.Series(pred_y))

		plt.title(title)
		sb.heatmap(confusion_matrix(y_true, pred_y), cmap='PuBu', annot=True, fmt='g', cbar=False)
		plt.xlabel(f'predict class')
		plt.ylabel('actual class')
		plt.xticks((0.5, 1.5, 2.5), ('1-8', '9-11', '12-29'))
		plt.yticks((0.5, 1.5, 2.5), ('1-8', '9-11', '12-29'))
		ax = plt.gca()
		ax.xaxis.set_label_position('top')
		ax.xaxis.set_ticks_position('top')
		plt.tight_layout()
		plt.show()

